from kuai.sim import *
from kuai.mol.io import read_molfile
from kuai.sim.dff2kuai import read_ppffile
from kuai.unit import default_units
from numpy import zeros, double

index, parameters = read_ppffile('C2H5OH and Acetone.ppf')

mol1 = read_molfile('C2H5OH.msd')
mol2 = read_molfile('Acetone.msd')

model = setup_model([(mol1, 100), (mol2, 100)], index, parameters)

A, _, _ = default_units.parse(1, 'A')

coords = []
with open('C2H5OH and Acetone.msd') as file:
    for i in file:
        if i.startswith('PBC:'):
            model.pbc = PBC.parse(i.strip())
            model.pbc *= A
        else:
            try:
                natoms = int(i.strip())
                break
            except ValueError:
                continue
    assert natoms == len(mol1.atoms) * 100 + len(mol2.atoms) * 100
    for i in range(natoms):
        tokens = file.next().split()
        assert len(tokens) == 12
        coords.append(float(tokens[6])*A)
        coords.append(float(tokens[7])*A)
        coords.append(float(tokens[8])*A)
model.coords = coords
# del model.charges

funcs = make_efunc(model, index)

for i in funcs:
    if i.name == 'pair-nonbond/lj12_6':
        i.cutoff = 10 * A       # Cutoff = 10 A

fe, _, _ = default_units.format(1, 'kcal/mol')
fg, _, _ = default_units.format(1, 'kcal/mol A')

job = MolecularDynamicJob()

temperature_c = BasicTemperatureControl()
pressure_c = BasicPressureControl()

controls = [temperature_c, pressure_c, SimpleReporter()]
for i in controls:
    i.period = 1000

model.speeds = zeros(len(coords), double)
model.potential_e, model.grediants = job.getEG(model, funcs)

step_length, _, _ = default_units.parse(1, 'fs')

from time import clock

# job.setCPU(4)

t1 = clock()
# job.run_md(model, funcs, controls, 10000)
job.run_md(model, funcs, controls, 10000, step_length)
print clock()-t1
print sum(model.potential_e) * fe

print "\n\n\n\n\n\n"

fg, _, _ = default_units.format(1, 'N/mol')
print model.grediants * fg
